import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.layers import Conv2D,Dense,Lambda,Activation,MaxPool2D,GlobalAveragePooling2D,Add,AvgPool2D,BatchNormalization
from tensorflow.keras.models import Model

class ConvModule():

    def __init__(self,in_channels,out_channels,kernel_size,stride,padding):
        self.Conv = Conv2D(out_channels,kernel_size,stride,padding)
        self.BatchNorm = BatchNormalization()
        self.activation = Activation(tf.nn.relu)

    def __call__(self, x):
        out = self.Conv(x)
        out = self.BatchNorm(out)
        out = self.activation(out)
        return out

# class BasicBlock():
#     def __init__(self):
#         pass

class Bottleneck():
    expansion = 4
    def __init__(self,inplance,plance,stride,down_sample):
        self.conv1_stride = 1
        self.conv2_stride = stride
        self.down_sample = down_sample
        self.CNR1 = ConvModule(inplance,plance,1,(1,1),'valid')
        self.CNR2 = ConvModule(plance,plance,3,(stride,stride),'same')
        self.Conv3 = Conv2D(plance*self.expansion,1,(1,1),'valid')
        self.BatchNorm = BatchNormalization()

    def __call__(self, x):
        short_cut = x
        if self.down_sample is not None:
            short_cut = self.down_sample(short_cut)
        out = self.CNR1(x)
        out = self.CNR2(out)
        out = self.Conv3(out)
        out = self.BatchNorm(out)

        out = Add()([out,short_cut])
        return out


class ResLayer():
    def __init__(self,inplance,plance,stride,block,block_num,avg_down,down_sample_first=True,**kwargs):
        """
        这里的kwargs是为了继承复写block模块是传入新增参数用的。
        :param inplance:
        :param plance:
        :param stride:
        :param block:
        :param block_num:
        :param avg_down:
        :param down_sample_first:
        :param kwargs:
        """
        down_sample = None
        if stride != 1 or inplance != plance*block.expansion:
            down_sample = self._make_down_sample_layer(stride,avg_down,plance*block.expansion)

        self.blocks = []
        if down_sample_first:
            self.blocks.append(block(inplance=inplance,plance=plance,stride=stride,down_sample=down_sample,**kwargs))
            inplance = plance * block.expansion
            for _ in range(block_num-1):
                self.blocks.append(block(inplance=inplance,plance=plance,stride=1,down_sample=None,**kwargs))
        else:
            inplance = plance * block.expansion
            for _ in range(block_num - 1):
                self.blocks.append(block(inplance=inplance, plance=inplance, stride=1, down_sample=None, **kwargs))
            self.blocks.append(block(inplance=inplance, plance=plance, stride=stride, down_sample=down_sample, **kwargs))

    def _make_down_sample_layer(self,stride,avg_down,out_channels):
        def down_sample(x):
            out = x
            conv_stride = stride
            if avg_down:
                out = AvgPool2D(pool_size=(stride,stride),padding='valid')(out)
                conv_stride = 1
            out = Conv2D(out_channels,1,(conv_stride,conv_stride),'valid')(out)
            out = BatchNormalization()(out)
            return out
        return down_sample

    def __call__(self, x):
        out = x
        for block in self.blocks:
            out = block(out)
        return out


class ResNet():
    arch_settings = {
        # 18: (BasicBlock, (2, 2, 2, 2)), #todo:add basic block
        # 34: (BasicBlock, (3, 4, 6, 3)),
        50: (Bottleneck, (3, 4, 6, 3)),
        101: (Bottleneck, (3, 4, 23, 3)),
        152: (Bottleneck, (3, 8, 36, 3))
    }
    def __init__(self,depth,in_channels,stem_channels=None,base_channels=64,num_stages=4,strides=(1,2,2,2),deep_stem=True,avg_down=False):
        self.block,stage_blocks = self.arch_settings[depth]
        self.stage_blocks = stage_blocks[:num_stages]
        if stem_channels is None:
            self.stem_channels = base_channels
        else:
            self.stem_channels = stem_channels
        self.base_channels = base_channels
        self.strides = strides
        self.deep_stem = deep_stem
        self.avg_down = avg_down

        self.stem_layer = self._make_stem_layer(in_channels,self.stem_channels)

        in_plance = self.stem_channels
        self.reslayers = []
        for i,num_block in enumerate(self.stage_blocks):
            stride = strides[i]
            plances = self.stem_channels * 2**i
            self.reslayers.append(self.make_res_layer(inplance=in_plance,plance=plances,stride=stride,block=self.block,block_num=num_block,avg_down=self.avg_down))
            in_plance = self.block.expansion * plances

    def _make_stem_layer(self,in_channels,stem_channels):
        def stem_layer(x):
            out = ConvModule(in_channels,stem_channels//2,3,(2,2),'same')(x)
            out = ConvModule(stem_channels//2,stem_channels//2,3,(1,1),'same')(out)
            out = ConvModule(stem_channels//2,stem_channels,3,(1,1,),'same')(out)
            return out
        return stem_layer

    def make_res_layer(self,**kwargs):
        return ResLayer(**kwargs)


    def __call__(self,x):
        out = self.stem_layer(x)
        out = MaxPool2D(pool_size=3,strides=2,padding='same')(out)
        for res_layer in self.reslayers:
            out = res_layer(out)
        return out

class ResNetV1d(ResNet):
    def __init__(self,**kwargs):
        super(ResNetV1d,self).__init__(in_channels=3,deep_stem=True,avg_down=True,**kwargs)

def create_model(input, clase_num):
    x = ResNetV1d(depth=50)(input)
    x = GlobalAveragePooling2D(name='avg_pool')(x)
    fc_out = Dense(clase_num, kernel_initializer="he_normal", use_bias=False, name="fc_NObias",activation='softmax')(x)
    model = Model(inputs=input, outputs=fc_out)
    return model

